#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：broadview_coss 
@File    ：ssh_client.py
@IDE     ：PyCharm 
@Author  ：kdly
@Date    ：2025/2/10 下午2:16 
'''
import re
import math
import time
import random
import string
import select
import hashlib

import logging

from django.core.cache import cache
from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
from paramiko.client import SSHClient, AutoAddPolicy

from hosts_mgr.models import CmdRecord
from utils.encrypt import Password

logger = logging.getLogger()
logger.setLevel(logging.INFO)  # 修改权限


class LocalSSHClient:
    def __init__(self, ip_addr, port, username, password, key_filename=''):
        """
        初始化 SSH 客户端
        :param ip_addr: 服务器地址
        :param port: 端口号，默认 22
        :param username: 用户名
        :param password: 密码（可选）
        :param key_filename: 私钥文件路径（可选）
        """
        self.ip_addr = ip_addr
        self.port = port
        self.username = username
        self.password = password
        self.key_filename = key_filename
        self.client = self.__connect()
        self.__encryption_pass = Password()

    def __connect(self):
        """
        连接到服务器
        """
        try:
            client = SSHClient()
            client.set_missing_host_key_policy(AutoAddPolicy())  # 自动添加主机密钥
            logger.info(f"尝试连接到 {self.ip_addr}:{self.port}...")
            if self.key_filename:
                # 使用私钥认证
                client.connect(
                    hostname=self.ip_addr,
                    port=self.port,
                    username=self.username,
                    key_filename=self.key_filename,
                    timeout=5,
                )
            else:
                # 使用密码认证
                client.connect(
                    hostname=self.ip_addr,
                    port=self.port,
                    username=self.username,
                    password=self.password,
                    timeout=10,
                )
            logger.info(f"SSH 成功连接到 {self.ip_addr}")
            return client
        except Exception as e:
            logger.error(f"SSH 连接失败: {e}")
            raise

    def execute_command_ws(self, server_type, command, host_id, hostname, user, username):
        """带重试机制的入口方法"""
        max_retries = 1
        retry_delay = 1  # 秒
        last_exception = ''
        try:
            for attempt in range(max_retries):
                logger.info(f"尝试第 {attempt + 1} 次执行命令...")
                if server_type == 1:
                    return self.__execute_command_ws_core(command, host_id, hostname, user)
                else:
                    # 生成唯一临时文件名
                    charset = string.ascii_letters + string.digits
                    random_str = ''.join(random.choice(charset) for _ in range(32))
                    temp_script = fr"C:\Users\{username}\Downloads\shell_script_{random_str}.bat"

                    channel_layer = get_channel_layer('default')
                    md5 = hashlib.md5()
                    md5.update(f"job_{host_id}_{user}".encode('utf-8'))
                    md5_group = md5.hexdigest()

                    cmd = fr"""@echo off
                    chcp 65001 > nul
                    {command}
                    set EXIT_CODE=%ERRORLEVEL%
                    exit /b %EXIT_CODE%""".encode("gbk")
                    # 使用 SFTP 传输文件
                    sftp = self.client.open_sftp()

                    # 将内容写入远程文件
                    with sftp.file(temp_script, 'wb') as remote_file:
                        remote_file.write(cmd)

                    stdin, stdout, stderr = self.client.exec_command(fr"call {temp_script}")

                    output = stdout.read().decode('utf-8')
                    error = stderr.read().decode('utf-8')

                    exit_status = stdout.channel.recv_exit_status()
                    self.client.exec_command(f"del {temp_script}")
                    async_to_sync(channel_layer.group_send)(
                        md5_group,
                        {"type": "send.message",
                         "message": {"code": 200, "type": "job_msg", "data": output}}
                    )
                    if error:
                        err = f"windows命令执行错误: {error}"
                        logger.error(err)
                        async_to_sync(channel_layer.group_send)(
                            md5_group,
                            {"type": "send.message",
                             "message": {"code": 500, "type": "job_error", "data": f"\x1B[31m错误: {error}\x1B[0m"}}
                        )
                        self.__exit_and_save_to_db(exit_status, hostname, channel_layer, md5_group, command, user)
                        raise err
                    self.__exit_and_save_to_db(exit_status, hostname, channel_layer, md5_group, command, user)
                    sftp.close()
                    self.client.close()
        except Exception as e:
            last_exception = e
            logger.warning(f"命令执行失败，第 {attempt + 1} 次重试，错误: {str(e)}")
            self.close(hostname)
            if attempt < max_retries - 1:
                time.sleep(retry_delay)
            else:
                logger.error(f"命令执行达到最大重试次数 {max_retries}")
                raise f"命令执行失败，重试{max_retries}次后仍失败: {str(last_exception)}"

    def __execute_command_ws_core(self, command, host_id, hostname, user):
        """实际执行逻辑"""
        logger.info(f"准备在主机 {hostname} 执行命令: '{command}'")

        if not self.client:
            self.connect()

        try:
            transport = self.client.get_transport()
            channel = transport.open_session()
            channel.get_pty(term='xterm-256color', width=160, height=50)

            # 生成唯一临时文件名
            charset = string.ascii_letters + string.digits
            random_str = ''.join(random.choice(charset) for _ in range(32))
            temp_script = f"/tmp/shell_script_{random_str}.sh"

            setup_script = (f"cat >{temp_script}<<'EOF'\n{command}\nEOF\n"
                            f"chmod +x {temp_script}\n"
                            f"bash -i {temp_script}\n"
                            f"exit_status=$?\n"
                            f"rm -rf {temp_script}\n"
                            f"exit $exit_status")

            channel.exec_command(setup_script)

            logger.info(f"命令 '{command}' 已发送到主机 {hostname}")

            # 非阻塞模式
            channel.setblocking(0)

            # 创建epoll对象
            epoll = select.epoll()
            epoll.register(channel.fileno(), select.EPOLLIN | select.EPOLLERR | select.EPOLLHUP)

            # WebSocket配置
            channel_layer = get_channel_layer('default')
            md5 = hashlib.md5()
            md5.update(f"job_{host_id}_{user}".encode('utf-8'))
            md5_group = md5.hexdigest()

            # 发送开始消息
            async_to_sync(channel_layer.group_send)(
                md5_group,
                {"type": "send.message",
                 "message": {"code": 200, "type": "job_msg", "data": '\x1B[32m开始执行脚本...\x1B[0m\r\n\r\n'}}
            )

            time.sleep(1)  # 确保前端准备好
            last_data_time = time.time()  # 最后数据接收时间
            timeout = 5  # 超时时间

            while True:
                # 检查前端连接状态
                if not cache.get(f'group_activity:{md5_group}'):
                    logger.warning(f"前端断开，终止连接：{hostname}")
                    self.__cleanup(epoll, channel, hostname, md5_group)
                    return

                # 检查数据超时
                if time.time() - last_data_time > timeout:
                    logger.warning(f"{hostname} 超过{timeout}秒无数据")
                    async_to_sync(channel_layer.group_send)(
                        md5_group,
                        {"type": "send.message",
                         "message": {"code": 408, "type": "job_error",
                                     "data": f"\r\n\x1B[31m错误: 超过{timeout}秒无响应，连接已断开\x1B[0m\r\n"}}
                    )
                    self.__cleanup(epoll, channel, hostname, md5_group)
                    return

                # 等待事件（1秒超时）
                events = epoll.poll(1000)
                data_received = False

                for fileno, event in events:
                    if fileno != channel.fileno():
                        continue

                    # 处理错误事件
                    if event & (select.EPOLLERR | select.EPOLLHUP):
                        logger.error("SSH通道错误")
                        self.__cleanup(epoll, channel, hostname, md5_group)
                        return

                    # 处理数据事件
                    if event & select.EPOLLIN:
                        try:
                            # 处理标准输出
                            if channel.recv_ready():
                                chunk = channel.recv(8192).decode('utf-8', 'ignore')
                                if chunk:
                                    data_received = True
                                    time.sleep(0.1)
                                    async_to_sync(channel_layer.group_send)(
                                        md5_group,
                                        {"type": "send.message",
                                         "message": {"code": 200, "type": "job_msg", "data": chunk}}
                                    )

                            # 处理错误输出
                            if channel.recv_stderr_ready():
                                chunk_err = channel.recv_stderr(4096).decode('utf-8', 'ignore')
                                if chunk_err:
                                    data_received = True
                                    async_to_sync(channel_layer.group_send)(
                                        md5_group,
                                        {"type": "send.message",
                                         "message": {"code": 500, "type": "job_error",
                                                     "data": f"\x1B[31m错误: {chunk_err}\x1B[0m"}}
                                    )

                        except Exception as e:
                            logger.error(f"数据读取失败: {e}")
                            self.__cleanup(epoll, channel, hostname, md5_group)
                            return

                # 更新最后数据接收时间
                if data_received:
                    last_data_time = time.time()

                # 检查命令是否完成
                if channel.exit_status_ready():
                    # 读取剩余输出
                    while channel.recv_ready():
                        chunk = channel.recv(4096).decode('utf-8', 'ignore')
                        async_to_sync(channel_layer.group_send)(
                            md5_group,
                            {"type": "send.message", "message": {"code": 200, "type": "job_msg", "data": chunk}}
                        )
                    break

                time.sleep(0.1)

            # 清理资源
            self.__cleanup(epoll, channel, hostname, md5_group)

            # 记录执行结果
            exit_status = channel.recv_exit_status()
            self.__exit_and_save_to_db(exit_status, hostname, channel_layer, md5_group, command, user)

        except Exception as e:
            logger.error(f"执行出错: {e}")
            self.close(hostname)
            raise

    def __exit_and_save_to_db(self, exit_status, hostname, channel_layer, md5_group, command, user):
        code_str = f'\r\n\x1B[31m执行结束，退出状态: {exit_status}\x1B[0m\r\n'
        if exit_status == 0:
            code_str = f'\r\n\x1B[32m执行成功，{hostname} 退出状态: {exit_status}\x1B[0m\r\n'

        # 发送结束通知
        for _ in range(3):
            if exit_status == 0:
                async_to_sync(channel_layer.group_send)(
                    md5_group,
                    {"type": "send.message",
                     "message": {"code": 200, "type": "job_close", "data": code_str}}
                )
            else:
                async_to_sync(channel_layer.group_send)(
                    md5_group,
                    {"type": "send.message",
                     "message": {"code": 500, "type": "job_error", "data": code_str}}
                )

            time.sleep(0.5)

        # 数据库记录（带重试）
        max_db_retries = 5
        for attempt in range(max_db_retries):
            try:
                if exit_status == 0:
                    CmdRecord.objects.create(cmd_content=command, hostname=hostname, user=user, status=1)
                else:
                    CmdRecord.objects.create(cmd_content=command, hostname=hostname, user=user, status=0)
                break
            except Exception as e:
                if attempt == max_db_retries - 1:
                    logger.error(f"数据库记录失败: {e}")
                time.sleep(1)

    def __cleanup(self, epoll, channel, hostname, group_key):
        """资源清理"""
        try:
            epoll.unregister(channel.fileno())
            epoll.close()
        except:
            pass
        self.close(hostname)
        cache.delete(f'group_activity:{group_key}')

    def execute_command_progress(self, dest_path, dest_host_obj, src_path, src_host_obj):
        """
        执行命令获取进度值
        :param dest_path: 文件名
        :param dest_host_obj: 目标主机对象
        :param src_path: 源地址
        :param src_host_obj: 源主机对象
        :return: 进度值
        """
        try:
            # 关键配置：使用伪终端获取实时输出
            transport = self.client.get_transport()
            channel = transport.open_session()
            channel.get_pty()  # 启用伪终端
            cmd = None
            file_size = 0
            if dest_host_obj.server_type == 1 and src_host_obj.server_type == 1:
                res = self.execute_command(f'du -s {src_path}')
                file_size = int(re.findall(r'(\d+)', res)[0]) * 1024

                cmd = (
                    f"echo 'put -r {src_path} {dest_path}' |sshpass -p '{self.__encryption_pass.decrypt(dest_host_obj.password)}' "
                    f"sftp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -vvv {dest_host_obj.user}@{dest_host_obj.connect_ip}"
                )
            elif dest_host_obj.server_type == 2 and src_host_obj.server_type == 1:
                res = self.execute_command(f'du -s {src_path}')
                file_size = int(re.findall(r'(\d+)', res)[0]) * 1024

                # 使用正则表达式替换单个反斜杠为双反斜杠
                # 正则表达式 r'(?<!\\)\\(?!\\)' 匹配不与其他反斜杠相邻的单个反斜杠
                new_path = re.sub(r"\\", "/", dest_path)
                cmd = (
                    f"echo 'put -r {src_path} /{new_path}' |sshpass -p '{self.__encryption_pass.decrypt(dest_host_obj.password)}' "
                    f"sftp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -vvv {dest_host_obj.user}@{dest_host_obj.connect_ip}"
                )

            elif src_host_obj.server_type == 2 and dest_host_obj.server_type == 1:
                password = self.__encryption_pass.decrypt(src_host_obj.password)
                client = LocalSSHClient(ip_addr=src_host_obj.connect_ip, port=src_host_obj.port,
                                        username=src_host_obj.user, password=password)
                res = client.execute_command(
                    fr'powershell -command "Get-ChildItem -Path {src_path} -Recurse | Measure-Object -Property Length -Sum | Select-Object Sum"',
                    win=1)
                file_size = int(re.findall(r'(\d+)', res)[0])
                new_path = re.sub(r"\\", "/", src_path)
                cmd = (
                    f"echo 'get -r /{new_path} {dest_path}' |sshpass -p '{self.__encryption_pass.decrypt(src_host_obj.password)}' "
                    f"sftp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -vvv {src_host_obj.user}@{src_host_obj.connect_ip}"
                )

            channel.exec_command(cmd)

            # 获取Channel层
            channel_layer = get_channel_layer('default')

            last_progress = 0

            # 获取前端ws连接过来的 file_name + host_id 的md5值
            md5 = hashlib.md5()
            md5.update(f"{dest_path}{dest_host_obj.id}".encode('utf-8'))
            md5_group = md5.hexdigest()
            buffer = ""  # 初始化缓冲区
            MAX_BUFFER_SIZE = 65536  # 64KB 最大缓冲区限制

            transport_size = 0
            size_num = 0
            while not channel.exit_status_ready():
                while channel.recv_ready():
                    time.sleep(0.02)
                    # 接收数据并添加到缓冲区
                    chunk = channel.recv(1024).decode('utf-8', 'ignore')
                    buffer += chunk

                    # 安全机制：防止缓冲区无限增长
                    if len(buffer) > MAX_BUFFER_SIZE:
                        # 丢弃旧数据
                        buffer = buffer[-MAX_BUFFER_SIZE // 2:]  # 保留最新的一半数据

                    # 处理缓冲区中的完整行
                    while '\n' in buffer:
                        line, buffer = buffer.split('\n', 1)  # 分割出第一行
                        line = line.strip()

                        # 原有的匹配和处理逻辑
                        if src_host_obj.server_type == 2:
                            match = re.findall(r'Received data (\d+) -> (\d+)', line)
                            if match:
                                match = match[0]
                                match = tuple([int(match[1]) - int(match[0])])
                        else:
                            match = re.findall(r'In write loop, ack for \d+ (\d+) bytes at', line)
                        if match:
                            transport_size += int(match[0])
                            if not transport_size:
                                continue
                            percent = int(transport_size / file_size * 100)

                            if size_num == percent:
                                continue

                            async_to_sync(channel_layer.group_send)(
                                md5_group,
                                {"type": "send.message", "message": {"code": 200, "type": "progress", "data": percent}}
                            )
                            size_num = last_progress = percent

                            logger.info(
                                f"{src_host_obj.hostname} ==> {dest_host_obj.hostname}:{dest_path} Progress: {percent}"
                            )
            # 强制发送最终进度
            if last_progress < 100:
                async_to_sync(channel_layer.group_send)(
                    md5_group,  # 组名需与 WebSocket 路由中的 group_name 一致
                    {
                        "type": "send.message",  # 对应 SSHConsumer 中的方法名（下划线式命名）
                        "message": {
                            "code": 200,
                            "type": "progress",
                            "data": 100  # 发送实际进度而非分块进度
                        }
                    }
                )
        except Exception as e:
            logger.error(f"传输失败: {e}")
            raise
        finally:
            self.close()

    def execute_command(self, command, status=0, win=0):
        """
        执行命令
        :param command: 要执行的命令
        :param status: 是否返回状态码, 1：是，0：否
        :param win: 是否是win，1：是
        :return: 命令输出
        """
        try:
            stdin, stdout, stderr = self.client.exec_command(command)
            if win == 1:
                output = stdout.read().decode(encoding="gbk", errors="strict")
                error = stderr.read().decode(encoding="gbk", errors="strict")
            else:
                output = stdout.read().decode('utf-8')
                error = stderr.read().decode('utf-8')
            if error and stdout.channel.recv_exit_status():
                logger.error(f"命令执行错误: {error}")
                # 将错误信息附加到异常中
                raise Exception(error)
            if status:
                return stdout.channel.recv_exit_status()  # 获取命令的退出状态码
            return output
        except Exception as e:
            raise e

    def upload_file(self, file_path, request):
        try:
            # 打开 SFTP 会话
            sftp = self.client.open_sftp()
            file_obj = request.FILES.get("file")
            # 上传文件
            sftp.put(file_obj, file_path)
            sftp.close()
            logger.info("文件上传成功")
            return True
        except Exception as e:
            logger.error(f"文件上传失败{e}")
            raise

    def get_all_msg(self):
        try:
            msg_dict = {}
            res = self.execute_command(command='cat /etc/os-release')
            os_release = re.findall('PRETTY_NAME="(.*)"', res)[0]
            # 如果 /etc/os-release 中没有找到，尝试使用 lsb_release 命令
            if not os_release:
                lsb_output = self.execute_command(command='lsb_release -d')
                if lsb_output.startswith("Description:"):
                    os_release = lsb_output.split(":")[1].strip()
            msg_dict['os_release'] = os_release
            mem_res = self.execute_command(command='cat /proc/meminfo|grep MemTotal')
            mem_info = mem_res.split()
            mem_info = math.ceil(int(mem_info[1]) / 1024 / 1024)
            cpu_info = self.execute_command(command='grep -c "^processor" /proc/cpuinfo')
            msg_dict['os_conf'] = f"{cpu_info.strip()}C{mem_info}G"
            res = self.execute_command(command='hostname -I')
            for ip in res.strip().split():
                if len(ip) > 15:  # ignore ipv6
                    continue
                msg_dict['private_ip'] = ip
                break
            res = self.execute_command(command='curl http://ifconfig.me 2> /dev/null')
            msg_dict['public_ip'] = res
            return msg_dict
        except Exception as e:
            logger.error(f"get_all_msg 失败{e}")
            raise

    def get_win_msg(self):
        try:
            msg_dict = {}
            res = self.execute_command(command='systeminfo', win=1)
            # 提取OS名称
            os_name_match = re.search(r"OS 名称:\s+(.+)", res)
            msg_dict['os_release'] = os_name_match.group(1).strip() if os_name_match else "Unknown OS"

            # 提取内存大小并转换为GB
            memory_match = re.search(r"物理内存总量:\s+([\d,]+) MB", res)
            if memory_match:
                memory_mb = int(memory_match.group(1).replace(",", ""))
                memory_gb = memory_mb // 1024
            else:
                memory_gb = 0

            # 提取CPU核数
            cpu_count_match = re.search(r"处理器:\s+安装了 (\d+) 个处理器", res)
            cpu_info = cpu_count_match.group(1) if cpu_count_match else "0"

            msg_dict['os_conf'] = f"{cpu_info.strip()}C{memory_gb}G"
            # 使用正则表达式匹配IPv4地址
            ip_pattern = r"IP 地址\s+\[01\]:\s+(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})"
            msg_dict['private_ip'] = re.search(ip_pattern, res).group(1)
            msg_dict['public_ip'] = self.execute_command(command='curl http://ifconfig.me', win=1)
            return msg_dict
        except Exception as e:
            logger.error(f"get_win_msg 失败{e}")
            raise

    def close(self, hostname=''):
        """
        关闭连接
        """
        if self.client:
            self.client.close()
            logger.warning(f"{hostname} SSH 连接已关闭")

    def __enter__(self):
        """支持 with 语句"""
        self.client = self.__connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持 with 语句"""
        self.close()


if __name__ == '__main__':
    # ssh_client = LocalSSHClient(ip_addr='127.0.0.1', port=22, username='debian', password='xxx')
    # # ssh_client.upload_file(file_name='/home/kdly/Pictures/55e9531d8492115707236aad49af7818.jpg')
    # print(ssh_client.execute_command(f"""{}"""))
    pass
